package com.sso.common.util;

import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import com.sso.common.entity.User;
import org.springframework.stereotype.Component;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.time.Instant;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * JWT 工具类
 * 用于生成和验证 JWT 令牌
 * 
 * @author SSO System
 * @since 2024-01-01
 */
@Component
public class JwtUtil {
    
    private static RSAKey rsaKey;
    private final JWSSigner signer;
    private final JWSVerifier verifier;
    
    /**
     * 构造函数，初始化 RSA 密钥对
     */
    public JwtUtil() {
        try {
            // 使用预生成的密钥对以提高启动性能
            this.rsaKey = generateOrLoadRsaKey();
            
            // 创建签名器和验证器
            this.signer = new RSASSASigner(rsaKey.toRSAPrivateKey());
            this.verifier = new RSASSAVerifier(rsaKey.toRSAPublicKey());
            
        } catch (Exception e) {
            throw new RuntimeException("初始化 JWT 工具类失败", e);
        }
    }
    
    /**
     * 生成或加载RSA密钥对
     * 为了提高性能，使用固定的密钥对
     */
    private RSAKey generateOrLoadRsaKey() throws Exception {
        // 在生产环境中，这些密钥应该从配置文件或密钥管理系统中加载
        // 这里为了演示使用固定的密钥对
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        
        return new RSAKey.Builder(publicKey)
                .privateKey(privateKey)
                .keyID("sso-key-id")
                .build();
    }
    
    /**
     * 生成访问令牌
     * 
     * @param username 用户名
     * @param authorities 用户权限
     * @param additionalClaims 额外的声明
     * @return JWT 令牌
     */
    public String generateAccessToken(String username, Set<String> authorities, Map<String, Object> additionalClaims) {
        try {
            Instant now = Instant.now();
            Instant expiresAt = now.plusSeconds(3600); // 1小时过期
            
            JWTClaimsSet.Builder claimsBuilder = new JWTClaimsSet.Builder()
                    .subject(username)
                    .issuer("sso-authorization-server")
                    .audience("sso-client")
                    .expirationTime(Date.from(expiresAt))
                    .notBeforeTime(Date.from(now))
                    .issueTime(Date.from(now))
                    .jwtID(java.util.UUID.randomUUID().toString());
            
            // 添加权限信息
            if (authorities != null && !authorities.isEmpty()) {
                claimsBuilder.claim("authorities", authorities);
            }
            
            // 添加额外声明
            if (additionalClaims != null) {
                additionalClaims.forEach(claimsBuilder::claim);
            }
            
            JWTClaimsSet claimsSet = claimsBuilder.build();
            
            SignedJWT signedJWT = new SignedJWT(
                    new JWSHeader.Builder(JWSAlgorithm.RS256)
                            .keyID(rsaKey.getKeyID())
                            .build(),
                    claimsSet
            );
            
            signedJWT.sign(signer);
            return signedJWT.serialize();
            
        } catch (Exception e) {
            throw new RuntimeException("生成访问令牌失败", e);
        }
    }
    
    /**
     * 生成刷新令牌
     * 
     * @param username 用户名
     * @return 刷新令牌
     */
    public String generateRefreshToken(String username) {
        try {
            Instant now = Instant.now();
            Instant expiresAt = now.plusSeconds(7 * 24 * 3600); // 7天过期
            
            JWTClaimsSet claimsSet = new JWTClaimsSet.Builder()
                    .subject(username)
                    .issuer("sso-authorization-server")
                    .audience("sso-client")
                    .expirationTime(Date.from(expiresAt))
                    .notBeforeTime(Date.from(now))
                    .issueTime(Date.from(now))
                    .jwtID(java.util.UUID.randomUUID().toString())
                    .claim("token_type", "refresh")
                    .build();
            
            SignedJWT signedJWT = new SignedJWT(
                    new JWSHeader.Builder(JWSAlgorithm.RS256)
                            .keyID(rsaKey.getKeyID())
                            .build(),
                    claimsSet
            );
            
            signedJWT.sign(signer);
            return signedJWT.serialize();
            
        } catch (Exception e) {
            throw new RuntimeException("生成刷新令牌失败", e);
        }
    }
    
    /**
     * 验证并解析 JWT 令牌
     * 
     * @param token JWT 令牌
     * @return JWT 声明集合
     */
    public JWTClaimsSet validateAndParseToken(String token) {
        try {
            SignedJWT signedJWT = SignedJWT.parse(token);
            
            // 验证签名
            if (!signedJWT.verify(verifier)) {
                throw new RuntimeException("JWT 签名验证失败");
            }
            
            JWTClaimsSet claimsSet = signedJWT.getJWTClaimsSet();
            
            // 验证过期时间
            if (claimsSet.getExpirationTime().before(new Date())) {
                throw new RuntimeException("JWT 令牌已过期");
            }
            
            // 验证生效时间
            if (claimsSet.getNotBeforeTime().after(new Date())) {
                throw new RuntimeException("JWT 令牌尚未生效");
            }
            
            return claimsSet;
            
        } catch (Exception e) {
            throw new RuntimeException("验证 JWT 令牌失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 从令牌中获取用户名
     * 
     * @param token JWT 令牌
     * @return 用户名
     */
    public String getUsernameFromToken(String token) {
        try {
            JWTClaimsSet claimsSet = validateAndParseToken(token);
            return claimsSet.getSubject();
        } catch (Exception e) {
            throw new RuntimeException("从令牌中获取用户名失败", e);
        }
    }
    
    /**
     * 检查令牌是否过期
     * 
     * @param token JWT 令牌
     * @return 是否过期
     */
    public boolean isTokenExpired(String token) {
        try {
            JWTClaimsSet claimsSet = validateAndParseToken(token);
            return claimsSet.getExpirationTime().before(new Date());
        } catch (Exception e) {
            return true; // 解析失败视为过期
        }
    }
    
    /**
     * 获取 RSA 公钥（用于资源服务器验证令牌）
     * 
     * @return RSA 公钥
     */
    public RSAKey getRsaKey() {
        return rsaKey.toPublicJWK();
    }
    
    /**
     * 为用户生成JWT令牌（用于单点登录）
     * 
     * @param user 用户对象
     * @return JWT令牌
     */
    public static String generateToken(User user) {
        try {
            // 创建RSA密钥对
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
            keyPairGenerator.initialize(2048);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            
            RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
            RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
            
            RSAKey rsaJWK = new RSAKey.Builder(publicKey)
                    .privateKey(privateKey)
                    .keyID("sso-portal-key")
                    .build();
            
            // 创建JWT签名器
            JWSSigner signer = new RSASSASigner(rsaJWK);
            
            // 设置JWT声明
            Instant now = Instant.now();
            Instant expiresAt = now.plusSeconds(24 * 3600); // 24小时过期
            
            JWTClaimsSet.Builder claimsBuilder = new JWTClaimsSet.Builder()
                    .subject(user.getUsername())
                    .issuer("portal-system")
                    .audience("sso-client")
                    .expirationTime(Date.from(expiresAt))
                    .notBeforeTime(Date.from(now))
                    .issueTime(Date.from(now))
                    .jwtID(java.util.UUID.randomUUID().toString());
            
            // 添加用户信息
            Map<String, Object> userInfo = new HashMap<>();
            userInfo.put("username", user.getUsername());
            if (user.getEmail() != null) {
                userInfo.put("email", user.getEmail());
                claimsBuilder.claim("email", user.getEmail());
            }
            if (user.getFullName() != null) {
                userInfo.put("fullName", user.getFullName());
                claimsBuilder.claim("fullName", user.getFullName());
            }
            
            // 添加权限信息
            if (user.getAuthorities() != null && !user.getAuthorities().isEmpty()) {
                claimsBuilder.claim("authorities", user.getAuthorities());
            }
            
            JWTClaimsSet claimsSet = claimsBuilder.build();
            
            // 创建签名的JWT
            SignedJWT signedJWT = new SignedJWT(
                    new JWSHeader.Builder(JWSAlgorithm.RS256)
                            .keyID(rsaJWK.getKeyID())
                            .build(),
                    claimsSet
            );
            
            // 签名JWT
            signedJWT.sign(signer);
            
            // 序列化JWT
            return signedJWT.serialize();
            
        } catch (Exception e) {
            throw new RuntimeException("生成单点登录令牌失败", e);
        }
    }
    
    /**
     * 验证单点登录令牌
     * 
     * @param token JWT令牌
     * @return 用户信息
     */
    public static Map<String, Object> validateSsoToken(String token) {
        try {
            SignedJWT signedJWT = SignedJWT.parse(token);
            JWTClaimsSet claimsSet = signedJWT.getJWTClaimsSet();
            
            // 验证过期时间
            if (claimsSet.getExpirationTime().before(new Date())) {
                throw new RuntimeException("SSO令牌已过期");
            }
            
            // 提取用户信息
            Map<String, Object> userInfo = new HashMap<>();
            userInfo.put("username", claimsSet.getSubject());
            userInfo.put("email", claimsSet.getClaim("email"));
            userInfo.put("fullName", claimsSet.getClaim("fullName"));
            userInfo.put("authorities", claimsSet.getClaim("authorities"));
            
            return userInfo;
        } catch (Exception e) {
            throw new RuntimeException("验证单点登录令牌失败: " + e.getMessage(), e);
        }
    }
}